package com.jboltai.capability.text;

import cn.hutool.core.util.StrUtil;
import com.jboltai.config.JBoltAIConfig;
import com.jboltai.event.EventAttrKey;
import com.jboltai.event.EventErrorType;
import com.jboltai.event.EventState;
import com.jboltai.event.text.ContentSourceType;
import com.jboltai.event.text.TextEvent;
import com.jboltai.event.text.TextExtractEvent;
import com.jboltai.event.text.TextSplitEvent;
import com.jboltai.util.share.AIEventThreadShareData;
import com.jboltai.util.text.TextUtil;
import org.apache.tika.exception.TikaException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.xml.sax.SAXException;

import java.io.File;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.concurrent.ConcurrentHashMap;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
import java.util.stream.Collectors;

public class TextEventProcessor implements Runnable{

    private static final Logger LOGGER = LoggerFactory.getLogger(JBoltAIConfig.LOGGER_NAME);


    private TextEvent event;


    public TextEventProcessor(TextEvent event) {
        this.event = event;

    }

    private static final ConcurrentHashMap<String, Pattern> patternMap = new ConcurrentHashMap<>();

    private static final String EMPTY_STR = "";


    @Override
    public void run() {
        if (this.event.getThreadShareData() != null) {
            AIEventThreadShareData.setShareData(this.event.getThreadShareData());
        }
        this.event.setAttr(EventAttrKey.EVENT_PROCESSOR, Thread.currentThread());
        try {
            if (event.setState(EventState.PROCESSING)) {
                switch (event.getType()) {
                    case TEXT_EXTRACT:
                        extractText();
                        break;
                    case TEXT_SPLIT:

                        splitText();
                        break;

                }
                event.setState(EventState.SUCCESS);

            }

        } catch (Throwable e) {
            if (e.getCause() != null) {
                e = e.getCause();
            }
            e.printStackTrace();
            if (e instanceof TikaException) {
                event.fail(EventErrorType.NOT_SUPPORT_DATA, "文件无法解析", e);
            } else if (e instanceof SAXException) {
                event.fail(EventErrorType.NOT_SUPPORT_DATA, "内容无法解析", e);
            } else {
                event.fail(EventErrorType.UNKNOWN, "未知错误，解析失败", e);
            }
        } finally {
            chunk = null;
            chunks = null;
            prevSplits = null;
            currentSplits = null;
            remainStr = null;
            event.complete();
        }

    }

    /**
     * 分割文本
     */
    private void splitText() {
        prevSplits = new ArrayList<>();
        currentSplits = new ArrayList<>();
        if (((TextSplitEvent)event).isStream()) {
            splitAsStream(((TextSplitEvent)event));
        } else {
            splitAsNormal(((TextSplitEvent)event));
        }
    }


    private SplitItem tempSplitItem = null;


    /**
     * 普通模式分割，先将文件的文本全部提取，再进行分割
     * @param event
     */
    private void splitAsNormal(TextSplitEvent event) {
        String content = event.getOriText();
        TextExtractResult extractResult = null;
        if (event.getFile() != null) {
            extractResult = TextUtil.extract(event.getFile(),event);
        } else if (event.getUrl() != null) {
            extractResult = TextUtil.extract(event.getUrl(),event);
        }
        if (extractResult != null) {
            content = extractResult.getContent();
        }
        if (StrUtil.isBlank(content)) {
            event.fail(EventErrorType.TEXT_EMPTY);
            return;
        }
        content = content.replaceAll("[ ]+", " ").replaceAll("\n+", "\n");

        if (StrUtil.isBlank(content)) {
            event.fail(EventErrorType.TEXT_EMPTY);
            return;
        }


        doSplit(content, event, false);
        if (chunk.length() > 0 || remainStr != null) {
            addChunk((chunk == null ? EMPTY_STR:chunk) + (remainStr == null ? EMPTY_STR : remainStr));
        }

        TextSplitResult textSplitResult = null;
        if (extractResult != null) {
            textSplitResult = new TextSplitResult(extractResult);
        } else {
            textSplitResult = new TextSplitResult(ContentSourceType.TEXT);
        }
        textSplitResult.setChuncks(chunks);
        //整体结束才调用success
        event.success(textSplitResult);
        event.setResult(textSplitResult);

    }

    private Pattern getPattern(String separators) {
        return patternMap.computeIfAbsent(separators, k -> Pattern.compile(separators));
    }


    /**
     * 流式分割，一边提取，一边分割
     * @param event
     */
    private void splitAsStream(TextSplitEvent event) {
        String content = event.getOriText();
        TextExtractResult extractResult = new TextExtractResult(null);
        event.setAttr(TextEvent.INNER_RESULT_KEY, extractResult);
        if (content != null) {
            extractResult.setContentSourceType(ContentSourceType.TEXT);
            doSplit(content, event, true);
        } else if (event.getFile() != null) {
            extractResult.setContentSourceType(ContentSourceType.LOCAL_FILE);
            extractResult.setSourceFile(event.getFile());
            extractResult.setTitle(event.getFile().getName());
            TextUtil.extractByChunk(event.getFile(), str -> {
                doSplit(str, event, true);
            },event);
        } else if (event.getUrl() != null) {
            extractResult.setSourceUrl(event.getUrl());
            TextUtil.extractByChunk(event.getUrl(), str -> {
                doSplit(str, event, true);
            },event);
        }
        //如果最后还留了点
        if (chunk.length() > 0 || remainStr != null) {
            event.success(new TextSplitResult(extractResult).setChuncks(Arrays.asList((chunk == null ? EMPTY_STR:chunk) + (remainStr == null ? EMPTY_STR : remainStr))));
        }
    }



    /**
     * 上一条的分块
     */
    private List<SplitItem> prevSplits = null;
    private List<SplitItem> currentSplits = null;

    private List<String> chunks = null;

    private String chunk = "";
    private double chunkToken = 0;
    String remainStr=null;

    /**
     * 执行分割
     * @param content
     * @param event
     * @param successOnChunk 是不是分一个块就执行一次success
     */
    private void doSplit(String content, TextSplitEvent event, boolean successOnChunk) {

        content = content.replaceAll("[ ]+", " ").replaceAll("\n+", "\n");
        if (remainStr != null) {
            content = remainStr + content;
            remainStr = null;
        }
        Pattern pattern = getPattern(Arrays.stream(event.getSeparators()).collect(Collectors.joining("|")));
        Matcher matcher = pattern.matcher(content);
        int lastIndex = 0;
        int start = -1, end = -1;

        while (matcher.find()) {

            start = matcher.start();
            end = matcher.end();

            tempSplitItem = new SplitItem(content.substring(lastIndex, start));
            processSplitItem(event,  successOnChunk);

            //拼接分割符
            tempSplitItem = new SplitItem(content.substring(start, end));
            currentSplits.add(tempSplitItem);
            chunk += tempSplitItem.content;
            chunkToken += tempSplitItem.tokenCount;

            lastIndex = end;
        }
        if (lastIndex != content.length()) {
            remainStr = content.substring(lastIndex);

            //这部分不匹配分割符，所以不能执行分块
        }

    }

    /**
     * 处理分割条目
     * @param event
     * @param successOnChunk
     */
    private void processSplitItem(TextSplitEvent event, boolean successOnChunk) {
        // 如果当前块的token数量加上新的分割项的token数量超过了设定的块大小
        if (chunkToken + tempSplitItem.tokenCount > event.getChunkSize()) {
            // 超过限制了，需要立即进行分割
            // 如果是流式处理，那么立即返回成功
            if (successOnChunk) {
                TextExtractResult extractResult = event.getAttr(TextEvent.INNER_RESULT_KEY);
                event.success(new TextSplitResult(extractResult).setChuncks(Arrays.asList(chunk)));
            } else {
                // 否则，将当前块添加到块列表中
                addChunk(chunk);
            }
            // 清空当前块和token计数
            chunk = "";
            chunkToken = 0;
            // 清空prevSplits，并将prevSplits和currentSplits交换
            List<SplitItem> temp = currentSplits;
            prevSplits.clear();
            currentSplits = prevSplits;
            prevSplits = temp;
        }
        // 拼接重叠部分
        joinOverlap(event);
        // 将新的分割项添加到当前分割列表中
        currentSplits.add(tempSplitItem);
        // 更新当前块和token计数
        chunk += tempSplitItem.content;
        chunkToken += tempSplitItem.tokenCount;
    }


    /**
     * 拼接重叠部分
     * @param event
     */
    private void joinOverlap(TextSplitEvent event) {
        if (chunk.length() == 0 && prevSplits.isEmpty() == false && event.getChunkOverlap() > 0) {
            //一个新的分段，要从前一个拼接重叠部分
            double reserve = event.getChunkOverlap();
            int lastEndIndex = prevSplits.size() - 2;
            while (reserve > 0  && lastEndIndex >= 0) {
                //从上一个开始往前找，大小小于chunkOverlap的就追加在前面
                if (prevSplits.get(lastEndIndex).tokenCount < reserve ) {
                    chunk = prevSplits.get(lastEndIndex).content + prevSplits.get(lastEndIndex+1).content + chunk;
                    chunkToken += (prevSplits.get(lastEndIndex).tokenCount + prevSplits.get(lastEndIndex+1).tokenCount);
                    reserve -= (prevSplits.get(lastEndIndex).tokenCount + prevSplits.get(lastEndIndex+1).tokenCount);

                    lastEndIndex -=2;
                } else {
                    break;
                }
            }
            prevSplits.clear();
        }
    }




    /**
     * 解析文本
     */
    private void extractText() {
        if (((TextExtractEvent)event).isStream()) {
            extractAsStream(((TextExtractEvent)event));
        } else {
            extractAsNormal(((TextExtractEvent)event));
        }
    }


    /**
     * 流式提取文本，一边提取，一边触发success
     * @param event
     */
    private void extractAsStream(TextExtractEvent event) {
        if (event.getFile() != null) {
            TextUtil.extractByChunk(event.getFile(), str -> {
                File file = event.getFile();
                event.success(new TextExtractResult(ContentSourceType.LOCAL_FILE).setTitle(file.getName()).setSourceFile(file).setContent(str));
            }, event);
        } else if (event.getUrl() != null) {
            TextUtil.extractByChunk(event.getUrl(), str -> {
                TextExtractResult result = event.getAttr(TextEvent.INNER_RESULT_KEY);
                if (result != null) {
                    //如果这是个网页的话，只会触发一次success，并且在解析网页的时候，会生成一个TextResult，attr进event
                    event.success(result.setContent(str));
                } else {
                    event.success(new TextExtractResult(ContentSourceType.ONLINE_FILE).setSourceUrl(event.getUrl()).setContent(str).setTitle(event.getUrl().substring(event.getUrl().lastIndexOf(".") + 1)));
                }

            }, event);
        }
    }

    /**
     * 普通模式提取文本，一次性将文件的文本内容提取出来，触发success
     * @param event
     */
    private void extractAsNormal(TextExtractEvent event) {
        TextExtractResult result  = null;
        if (event.getFile() != null) {
            result = TextUtil.extract(event.getFile(), event);
        } else if (event.getUrl() != null) {
            result = TextUtil.extract(event.getUrl(), event);
        }
        event.setResult(result);
        if (StrUtil.isBlank(result.getContent())) {
            event.fail(EventErrorType.TEXT_EMPTY);
        } else {
            event.success(result);
        }
    }


    public void addChunk(String chunk) {
        if (chunks == null) {
            chunks = new ArrayList<>();
        }
        chunks.add(chunk);
    }


    /**
     * 分割后的条目
     */
    private static class SplitItem {
        /**
         * 内容
         */
        String content;
        /**
         * 该条目的token数，预估值
         */
        double tokenCount;

        public SplitItem(String content) {
            this.content = content;
            this.tokenCount = TextUtil.calToken(content);
        }
    }



}
